import logging
import time
import numpy as np
from PIL import Image
# 调整 ServerlessSpec 的导入路径（关键修改）
from pinecone import Pinecone, ServerlessSpec
from tqdm import tqdm
from tensorflow.keras.datasets import mnist

# 配置 logging（含日期）
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# -------------------------- 你的 Pinecone 配置（无需修改）--------------------------
PINECONE_API_KEY = "pcsk_2WNR5R_95XGmZAVXXQYSc2txi8pJwXVWQaUg2WmdnT1nREKVJRSXdTH4CaJ3Lq9NcfKwD3"
INDEX_NAME = "developer-quickstart-py"
VECTOR_DIM = 64  # 8x8 图像展平后维度（固定）
# -----------------------------------------------------------------------------------


def load_and_preprocess_mnist():
    """使用 TensorFlow 加载并预处理 MNIST 数据（缩放为 8x8，适配原逻辑）"""
    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()

    def resize_and_normalize(img):
        """将 28x28 原始图像 → 8x8 灰度图 → 归一化到 0-16"""
        # 转为 PIL 图像以便缩放
        img_pil = Image.fromarray(img)
        # 缩放到 8x8（与 Web 应用预处理一致）
        img_resized = img_pil.resize((8, 8), Image.Resampling.LANCZOS)
        # 颜色反转（黑底白字 → 白底黑字，匹配训练数据分布）
        img_inverted = Image.eval(img_resized, lambda x: 255 - x)
        # 归一化到 0-16（原代码逻辑，确保与推理时特征范围一致）
        return np.array(img_inverted, dtype=np.float32) / 16

    # 批量处理训练集和测试集
    train_images_8x8 = np.array([resize_and_normalize(img) for img in train_images]).reshape(-1, 64)
    test_images_8x8 = np.array([resize_and_normalize(img) for img in test_images]).reshape(-1, 64)

    # 按 8:2 划分数据：80% 用于上传 Pinecone，20% 用于测试准确率
    split_idx = int(len(train_images_8x8) * 0.8)
    pinecone_train_imgs = train_images_8x8[:split_idx]
    pinecone_train_labels = train_labels[:split_idx]
    test_imgs_subset = test_images_8x8[:int(len(test_images_8x8) * 0.2)]  # 取测试集20%，减少测试时间
    test_labels_subset = test_labels[:int(len(test_labels) * 0.2)]

    return pinecone_train_imgs, pinecone_train_labels, test_imgs_subset, test_labels_subset


def init_pinecone_index():
    """初始化 Pinecone 索引（适配最新库的类与方法）"""
    # 1. 实例化 Pinecone 客户端（替代旧版 pc.init()）
    pc_client = Pinecone(api_key=PINECONE_API_KEY)
    logging.info("✅ Pinecone 客户端初始化完成")

    # 2. 检查索引是否存在，不存在则创建
    if INDEX_NAME not in pc_client.list_indexes().names():
        logging.info(f"🔍 索引 '{INDEX_NAME}' 不存在，开始创建 Serverless 索引...")
        pc_client.create_index(
            name=INDEX_NAME,
            dimension=VECTOR_DIM,  # 必须与图像特征维度一致（64）
            metric="euclidean",    # KNN 手写数字识别常用欧氏距离
            spec=ServerlessSpec(   # 显式传入 Serverless 配置（关键修改）
                cloud="aws",
                region="us-east-1"
            )
        )
        # 等待索引创建完成（Serverless 索引创建需约10-15秒，避免后续连接失败）
        time.sleep(15)
        logging.info(f"✅ 索引 '{INDEX_NAME}' 创建完成")
    else:
        logging.info(f"✅ 索引 '{INDEX_NAME}' 已存在，直接连接")

    # 3. 返回已连接的索引对象
    return pc_client.Index(INDEX_NAME)


def upload_data_to_pinecone(index, images, labels):
    """上传预处理后的 MNIST 数据到 Pinecone（带进度条，适配最新 API）"""
    logging.info("📤 开始上传数据到 Pinecone...")
    vectors = []  # 存储批量向量数据（每批100条，避免请求过载）

    for idx, (img, label) in enumerate(zip(images, labels)):
        # 构造 Pinecone 要求的向量格式（id + values + metadata）
        vectors.append({
            "id": f"mnist-{idx}",          # 唯一ID（格式：mnist-0、mnist-1...）
            "values": img.tolist(),        # 64维特征向量（转为列表适配API）
            "metadata": {"label": int(label)}  # 存储数字标签，用于后续预测投票
        })

        # 每积累100条数据或到最后一条时，执行批量上传
        if len(vectors) >= 100 or idx == len(images) - 1:
            with tqdm(total=len(vectors), desc="📤 上传进度") as pbar:
                index.upsert(vectors=vectors)  # 执行上传
                pbar.update(len(vectors))      # 更新进度条
            vectors = []  # 清空批量列表，准备下一批

    # 打印上传结果
    upload_count = len(images)
    logging.info(f"✅ 成功上传 {upload_count} 条 MNIST 数据到 Pinecone")


def test_accuracy_with_k(index, test_images, test_labels, k=11):
    """测试 k=11 时的识别准确率（投票法，带进度条）"""
    logging.info(f"🧪 开始测试 k={k} 的准确率...")
    correct_count = 0  # 记录正确预测的样本数

    # 遍历测试集，逐样本预测
    with tqdm(total=len(test_images), desc="🧪 测试进度") as pbar:
        for img, true_label in zip(test_images, test_labels):
            # 向 Pinecone 发起查询：获取 top-k 个相似向量
            query_result = index.query(
                vector=img.tolist(),  # 测试样本的64维特征向量
                top_k=k,              # 取前11个相似结果（按需求固定k=11）
                include_metadata=True # 必须显式开启，才能获取元数据中的 label
            )

            # 投票法确定预测结果：统计 top-k 结果中出现次数最多的 label
            pred_labels = [int(match["metadata"]["label"]) for match in query_result["matches"]]
            pred_label = max(set(pred_labels), key=pred_labels.count)

            # 验证预测是否正确
            if pred_label == true_label:
                correct_count += 1

            pbar.update(1)  # 更新进度条

    # 计算并打印准确率
    accuracy = correct_count / len(test_images) if len(test_images) > 0 else 0.0
    logging.info(f"✅ k={k} 时的准确率为：{accuracy:.4f}（正确数/总数：{correct_count}/{len(test_images)}）")


if __name__ == "__main__":
    # 执行完整流程：加载数据 → 初始化索引 → 上传数据 → 测试准确率
    logging.info("🚀 开始执行 Pinecone MNIST 训练流程")
    
    # 1. 加载并预处理 MNIST 数据（TensorFlow 自动下载，无需手动操作）
    train_imgs, train_labels, test_imgs, test_labels = load_and_preprocess_mnist()
    logging.info(f"✅ MNIST 数据加载完成：训练集 {len(train_imgs)} 条，测试集 {len(test_imgs)} 条")
    
    # 2. 初始化 Pinecone 索引
    pinecone_index = init_pinecone_index()
    
    # 3. 上传训练数据到 Pinecone
    upload_data_to_pinecone(pinecone_index, train_imgs, train_labels)
    
    # 4. 测试 k=11 的准确率
    test_accuracy_with_k(pinecone_index, test_imgs, test_labels, k=11)
    
    logging.info("🎉 Pinecone MNIST 训练流程全部执行完成！")